{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JmL6bBy1yam6",
    "outputId": "95667fee-e18c-4fee-de93-90f16769369c"
   },
   "outputs": [],
   "source": [
    "######## Installations\n",
    "\n",
    "!pip install zep-cloud openai --upgrade"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lHIC512UV51K"
   },
   "outputs": [],
   "source": [
    "######## Imports\n",
    "\n",
    "import pandas as pd\n",
    "from openai import AsyncOpenAI\n",
    "from zep_cloud.client import AsyncZep\n",
    "from zep_cloud import Message, EntityEdge, EntityNode\n",
    "import os\n",
    "from dotenv import load_dotenv\n",
    "from pydantic import BaseModel, Field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "t-bBAEydL9su"
   },
   "outputs": [],
   "source": [
    "load_dotenv()\n",
    "\n",
    "ZEP_API_KEY = os.getenv(\"ZEP_API_KEY\")\n",
    "OPENAI_API_KEY = os.getenv(\"OPENAI_API_KEY\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gQN9V9UlXXR5"
   },
   "outputs": [],
   "source": [
    "######## Helper functions\n",
    "\n",
    "\n",
    "def nicely_display_one_data_sample(msc_dataset_df, idx):\n",
    "    \"\"\"\n",
    "    Given the MSC dataset's dataframe and an idx (between 0 and 499), print out one data sample\n",
    "    in a readable format.\n",
    "    \"\"\"\n",
    "\n",
    "    msc_dataset_df_keys = msc_dataset_df.keys()\n",
    "\n",
    "    for key in msc_dataset_df_keys:\n",
    "        keys_object = msc_dataset_df[key][idx]\n",
    "        print(f\"\\n*****{key}\")\n",
    "\n",
    "        # Handle metadata and self instruct differently\n",
    "        if key == \"metadata\":\n",
    "            continue\n",
    "        if key == \"self_instruct\":\n",
    "            print(\"QUESTION (B's Message):\", keys_object[\"B\"])\n",
    "            print(\"ANSWER (A's Message):\", keys_object[\"A\"])\n",
    "\n",
    "        for i in range(20):\n",
    "            try:\n",
    "                print(keys_object[i])\n",
    "            except:\n",
    "                break\n",
    "\n",
    "\n",
    "def get_all_chat_sessions(msc_dataset_df, idx):\n",
    "    \"\"\"\n",
    "    Given the MSC dataset's dataframe and an idx (between 0 and 499), return a list of all\n",
    "    5 chat sessions in a nice format, only including the information needed for Zep.\n",
    "    \"\"\"\n",
    "    all_chat_sessions = []\n",
    "\n",
    "    # Gather the previous chat sessions\n",
    "    for i, prev_dialog in enumerate(msc_dataset_df[\"previous_dialogs\"][idx]):\n",
    "        prev_dialog_message_list = []\n",
    "        for msg_dict in prev_dialog[\"dialog\"]:\n",
    "            prev_dialog_message_list.append(msg_dict[\"text\"])\n",
    "        all_chat_sessions.append(\n",
    "            {\n",
    "                \"messages\": prev_dialog_message_list,\n",
    "                \"session_num\": i + 1,\n",
    "                \"time_num\": prev_dialog[\"time_num\"],\n",
    "                \"time_unit\": prev_dialog[\"time_unit\"],\n",
    "                \"time_back\": prev_dialog[\"time_back\"],\n",
    "            }\n",
    "        )\n",
    "\n",
    "    # Gather the most recent/newest chat session, always session #5\n",
    "    newest_dialog = msc_dataset_df[\"dialog\"][idx]\n",
    "    newest_dialog_message_list = []\n",
    "    for msg_dict in newest_dialog:\n",
    "        newest_dialog_message_list.append(msg_dict[\"text\"])\n",
    "    all_chat_sessions.append({\"messages\": newest_dialog_message_list, \"session_num\": 5})\n",
    "\n",
    "    return all_chat_sessions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######## Download the eval dataset from the official HuggingFace Source\n",
    "# Download this file and store it in data/msc.jsonl\n",
    "# https://huggingface.co/datasets/MemGPT/MSC-Self-Instruct\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_t5PZzC8M66d",
    "outputId": "d2ed3ca8-e245-4d9b-aa8f-accb550486ed"
   },
   "outputs": [],
   "source": [
    "######## Load the eval dataset\n",
    "\n",
    "msc_dataset_df = pd.read_json(\"locomo_eval/data/msc.jsonl\", lines=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "F3X8rQyVLRJf"
   },
   "outputs": [],
   "source": [
    "######## Start up Zep and OpenAI clients\n",
    "zep = AsyncZep(\n",
    "    api_key=os.getenv(\"ZEP_API_KEY\"),\n",
    "    base_url=\"https://api.development.getzep.com/api/v2\",\n",
    ")\n",
    "oai_client = AsyncOpenAI(\n",
    "    api_key=os.getenv(\"OPENAI_API_KEY\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 280
    },
    "id": "Y9Iy7VQ7wcap",
    "outputId": "77a6de57-427f-40e3-c579-eb0d206f0a72"
   },
   "outputs": [],
   "source": [
    "######## MSC Ingestion loop - ingest messages for each multi-session\n",
    "async def ingest_msc(zep: AsyncZep):\n",
    "    num_multi_sessions = 500\n",
    "\n",
    "    for multi_session_idx in range(num_multi_sessions):\n",
    "        if multi_session_idx < 0:\n",
    "            continue\n",
    "\n",
    "        # Create a unique Zep user and session for this pair of speakers\n",
    "        user_id = \"msc_experiment_user_\" + str(multi_session_idx)\n",
    "        session_id = \"msc_experiment_session_\" + str(multi_session_idx)\n",
    "\n",
    "        # Uncomment to delete existing users\n",
    "        # try:\n",
    "        #     await zep.user.delete(user_id)\n",
    "        #     await zep.memory.delete(session_id)\n",
    "        # except:\n",
    "        #     pass\n",
    "        #\n",
    "        # continue\n",
    "\n",
    "        await zep.user.add(user_id=user_id)\n",
    "        await zep.memory.add_session(\n",
    "            user_id=user_id,\n",
    "            session_id=session_id,\n",
    "        )\n",
    "\n",
    "        # Now add the messages as episodes to this user-session\n",
    "\n",
    "        all_chat_sessions = get_all_chat_sessions(msc_dataset_df, multi_session_idx)\n",
    "        for session_idx, session in enumerate(all_chat_sessions):\n",
    "            for msg_idx, msg in enumerate(session[\"messages\"]):\n",
    "                if msg_idx % 2 == 0:\n",
    "                    await zep.memory.add(\n",
    "                        session_id=session_id,\n",
    "                        messages=[Message(role=\"A\", role_type=\"norole\", content=msg)],\n",
    "                    )\n",
    "                else:\n",
    "                    await zep.memory.add(\n",
    "                        session_id=session_id,\n",
    "                        messages=[Message(role=\"B\", role_type=\"norole\", content=msg)],\n",
    "                    )\n",
    "\n",
    "\n",
    "await ingest_msc(zep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######## Define prompts for Deep Memory Retrieval (DMR) eval\n",
    "async def dmr_response(llm_client, context: str, question: str) -> str:\n",
    "    system_prompt = \"\"\"\n",
    "        You are speaker A and should respond to all questions in the the first person perspective of A\n",
    "        \"\"\"\n",
    "\n",
    "    prompt = f\"\"\"\n",
    "            Your task is to briefly answer the question. You are given the following context from the previous conversation. If you don't know how to answer the question, abstain from answering.\n",
    "                <CONTEXT>\n",
    "                {context}\n",
    "                </CONTEXT>\n",
    "                <QUESTION>\n",
    "                {question}\n",
    "                </QUESTION>\n",
    "\n",
    "            Respond with an ANSWER section containing your answer. As well as an EVIDENCE section containing the context that help you came to your conclusion, and an explanation of why that context is relevant.\n",
    "            \"\"\"\n",
    "\n",
    "    response = await llm_client.chat.completions.create(\n",
    "        model=\"gpt-4-mini\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": prompt},\n",
    "        ],\n",
    "        temperature=0,\n",
    "    )\n",
    "    result = response.choices[0].message.content or \"\"\n",
    "\n",
    "    return result\n",
    "\n",
    "\n",
    "class Grade(BaseModel):\n",
    "    is_correct: bool = Field(description=\"Whether or not the response is correct\")\n",
    "\n",
    "\n",
    "async def dmr_grader(\n",
    "    llm_client, question: str, gold_answer: str, response: str\n",
    ") -> bool:\n",
    "    system_prompt = \"\"\"\n",
    "        You are an expert grader that determines if answers to questions match a gold standard answer\n",
    "        \"\"\"\n",
    "\n",
    "    prompt = f\"\"\"         \n",
    "    I will give you a question, a correct answer, and a response from a model. Please answer true if the response contains the correct answer or touches on the same topic. Otherwise, answer false. If the response is equivalent to the correct answer or contains all the intermediate steps to get the correct answer, you should also answer true. If the response only contains a subset of the information required by the answer, answer false.\n",
    "            \n",
    "    <QUESTION>\n",
    "    B: {question}\n",
    "    </QUESTION>\n",
    "    <CORRECT ANSWER>\n",
    "    {gold_answer}\n",
    "    </CORRECT ANSWER>\n",
    "    <RESPONSE>\n",
    "    A: {response}\n",
    "    </RESPONSE>\n",
    "    \"\"\"\n",
    "\n",
    "    response = await llm_client.beta.chat.completions.parse(\n",
    "        model=\"gpt-4o-mini\",\n",
    "        messages=[\n",
    "            {\"role\": \"system\", \"content\": system_prompt},\n",
    "            {\"role\": \"user\", \"content\": prompt},\n",
    "        ],\n",
    "        response_format=Grade,\n",
    "        temperature=0,\n",
    "    )\n",
    "    result = response.choices[0].message.parsed\n",
    "\n",
    "    return result.is_correct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######## Baseline eval loop - for each multi-session, evaluate on summary and full context\n",
    "async def dmr_baseline():\n",
    "    num_multi_sessions = 500\n",
    "    eval_results = {\n",
    "        \"total\": num_multi_sessions,\n",
    "        \"full_conversation_correct\": 0,\n",
    "        \"full_conversation_accuracy\": 0,\n",
    "        \"conversation_summary_correct\": 0,\n",
    "        \"conversation_summary_accuracy\": 0,\n",
    "    }\n",
    "\n",
    "    for multi_session_idx in range(num_multi_sessions):\n",
    "        # Build contexts\n",
    "        full_conversation = \"\"\n",
    "\n",
    "        all_chat_sessions = get_all_chat_sessions(msc_dataset_df, multi_session_idx)\n",
    "        for session_idx, session in enumerate(all_chat_sessions):\n",
    "            for msg_idx, msg in enumerate(session[\"messages\"]):\n",
    "                if msg_idx % 2 == 0:\n",
    "                    full_conversation += \"A: \" + msg + \"\\n\"\n",
    "                else:\n",
    "                    full_conversation += \"B: \" + msg + \"\\n\"\n",
    "\n",
    "        speaker_a_summary_set = msc_dataset_df[\"summary_speaker_1\"][multi_session_idx]\n",
    "        speaker_b_summary_set = msc_dataset_df[\"summary_speaker_2\"][multi_session_idx]\n",
    "\n",
    "        speaker_a_summary = \"<SPEAKER A SUMMARY>\"\n",
    "        for session_summary in speaker_a_summary_set:\n",
    "            for fact in session_summary:\n",
    "                speaker_a_summary += \"\\n  \" + fact\n",
    "        speaker_a_summary += \"\\n</SPEAKER A SUMMARY>\"\n",
    "\n",
    "        speaker_b_summary = \"\\n<SPEAKER B SUMMARY>\"\n",
    "        for session_summary in speaker_b_summary_set:\n",
    "            for fact in session_summary:\n",
    "                speaker_b_summary += \"\\n  \" + fact\n",
    "        speaker_b_summary += \"\\n</SPEAKER B SUMMARY>\"\n",
    "\n",
    "        conversation_summary = speaker_a_summary + speaker_b_summary\n",
    "\n",
    "        # Set question and golden answer\n",
    "        question = msc_dataset_df[\"self_instruct\"][multi_session_idx][\"B\"]\n",
    "        gold_answer = msc_dataset_df[\"self_instruct\"][multi_session_idx][\"A\"]\n",
    "\n",
    "        # Prompt the LLM to answer the question\n",
    "        full_conversation_response = await dmr_response(\n",
    "            oai_client, full_conversation, question\n",
    "        )\n",
    "        conversation_summary_response = await dmr_response(\n",
    "            oai_client, conversation_summary, question\n",
    "        )\n",
    "\n",
    "        # print('QUESTION: ', question)\n",
    "        print(\"GOLDEN ANSWER\", gold_answer)\n",
    "        print(\"FULL CONVERSATION RESPONSE: \", full_conversation_response)\n",
    "        print(\"CONVERSATION SUMMARY RESPONSE: \", conversation_summary_response)\n",
    "\n",
    "        # Grade responses\n",
    "        full_conversation_grade = await dmr_grader(\n",
    "            oai_client, question, gold_answer, full_conversation_response\n",
    "        )\n",
    "        conversation_summary_grade = await dmr_grader(\n",
    "            oai_client, question, gold_answer, conversation_summary_response\n",
    "        )\n",
    "\n",
    "        if full_conversation_grade:\n",
    "            eval_results[\"full_conversation_correct\"] += 1\n",
    "        if conversation_summary_grade:\n",
    "            eval_results[\"conversation_summary_correct\"] += 1\n",
    "\n",
    "    eval_results[\"full_conversation_accuracy\"] = (\n",
    "        eval_results[\"full_conversation_correct\"] / num_multi_sessions\n",
    "    )\n",
    "    eval_results[\"conversation_summary_accuracy\"] = (\n",
    "        eval_results[\"conversation_summary_correct\"] / num_multi_sessions\n",
    "    )\n",
    "\n",
    "    print(eval_results)\n",
    "\n",
    "\n",
    "await dmr_baseline()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "######## Main eval loop - for each multi-session, evaluate Zep\n",
    "TEMPLATE = \"\"\"\n",
    "FACTS and ENTITIES represent relevant context to the current conversation.\n",
    "\n",
    "# These are the most relevant facts and their valid date ranges\n",
    "# format: FACT (Date range: from - to)\n",
    "<FACTS>\n",
    "{facts}\n",
    "</FACTS>\n",
    "\n",
    "# These are the most relevant entities\n",
    "# ENTITY_NAME: entity summary\n",
    "<ENTITIES>\n",
    "{entities}\n",
    "</ENTITIES>\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "def format_edge_date_range(edge: EntityEdge) -> str:\n",
    "    # return f\"{datetime(edge.valid_at).strftime('%Y-%m-%d %H:%M:%S') if edge.valid_at else 'date unknown'} - {(edge.invalid_at.strftime('%Y-%m-%d %H:%M:%S') if edge.invalid_at else 'present')}\"\n",
    "    return f\"{edge.valid_at if edge.valid_at else 'date unknown'} - {(edge.invalid_at if edge.invalid_at else 'present')}\"\n",
    "\n",
    "\n",
    "def compose_search_context(edges: list[EntityEdge], nodes: list[EntityNode]) -> str:\n",
    "    facts = [f\"  - {edge.fact} ({format_edge_date_range(edge)})\" for edge in edges]\n",
    "    entities = [f\"  - {node.name}: {node.summary}\" for node in nodes]\n",
    "    return TEMPLATE.format(facts=\"\\n\".join(facts), entities=\"\\n\".join(entities))\n",
    "\n",
    "\n",
    "async def dmr_eval_session(zep: AsyncZep, idx: int) -> bool:\n",
    "    # Set user values\n",
    "    user_id = \"msc_experiment_user_\" + str(idx)\n",
    "    session_id = \"msc_experiment_session_\" + str(idx)\n",
    "\n",
    "    # Now we want to prompt an LLM augmented with Zep memory to answer the question\n",
    "    question = msc_dataset_df[\"self_instruct\"][idx][\"B\"]\n",
    "    gold_answer = msc_dataset_df[\"self_instruct\"][idx][\"A\"]\n",
    "\n",
    "    # Get relevant facts and entities\n",
    "    edges_results = (\n",
    "        await zep.graph.search(\n",
    "            user_id=user_id,\n",
    "            reranker=\"cross_encoder\",\n",
    "            query=question,\n",
    "            scope=\"edges\",\n",
    "            limit=20,\n",
    "        )\n",
    "    ).edges\n",
    "    node_results = (\n",
    "        await zep.graph.search(\n",
    "            user_id=user_id, reranker=\"rrf\", query=question, scope=\"nodes\", limit=20\n",
    "        )\n",
    "    ).nodes\n",
    "\n",
    "    context = compose_search_context(edges_results, node_results)\n",
    "\n",
    "    # Prompt an LLM with relevant context\n",
    "    response = await dmr_response(oai_client, context, question)\n",
    "\n",
    "    # Grade responses\n",
    "    grade = await dmr_grader(oai_client, question, gold_answer, response)\n",
    "\n",
    "    if not grade:\n",
    "        print(\"IDX: \", idx)\n",
    "        print(\"CONTEXT: \", context)\n",
    "        print(\"QUESTION: \", question)\n",
    "        print(\"GOLDEN ANSWER: \", gold_answer)\n",
    "        print(\"RESPONSE: \", response)\n",
    "        print(\"GRADE: \", grade)\n",
    "\n",
    "    return grade\n",
    "\n",
    "\n",
    "async def dmr_eval(zep: AsyncZep):\n",
    "    num_multi_sessions = 500\n",
    "\n",
    "    eval_results = {\"correct\": 0, \"total\": num_multi_sessions, \"accuracy\": 0}\n",
    "\n",
    "    # grades = list(await asyncio.gather(*[dmr_eval_session(zep, multi_session_idx) for multi_session_idx in range(num_multi_sessions)]))\n",
    "    grades: list[bool] = []\n",
    "    for i in range(num_multi_sessions):\n",
    "        grades.append(await dmr_eval_session(zep, i))\n",
    "\n",
    "    for grade in grades:\n",
    "        if grade:\n",
    "            eval_results[\"correct\"] += 1\n",
    "\n",
    "        # Now that we have the results, delete the user\n",
    "        # await zep.user.delete(user_id)\n",
    "\n",
    "    eval_results[\"accuracy\"] = eval_results[\"correct\"] / num_multi_sessions\n",
    "\n",
    "    print(eval_results)\n",
    "\n",
    "\n",
    "await dmr_eval(zep)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
